import math
import pandas as pd
from collections import Counter


# 计算信息熵
def entropy(data):
    total = len(data)
    label_counts = Counter(data)
    ent = 0.0
    for count in label_counts.values():
        prob = count / total
        ent -= prob * math.log2(prob)
    return ent


# 计算信息增益
def info_gain(data, feature, target):
    total = len(data)
    feature_values = data[feature].unique()
    feature_info = 0.0
    for value in feature_values:
        subset = data[data[feature] == value]
        prob = len(subset) / total
        feature_info += prob * entropy(subset[target])
    return entropy(data[target]) - feature_info


# ID3算法
def id3(data, features, target):
    if len(data[target].unique()) == 1:
        return data[target].iloc[0]
    if not features:
        return Counter(data[target]).most_common(1)[0][0]

    best_feature = max(features, key=lambda f: info_gain(data, f, target))
    tree = {best_feature: {}}
    features = [f for f in features if f != best_feature]

    for value in data[best_feature].unique():
        subset = data[data[best_feature] == value]
        if subset.empty:
            tree[best_feature][value] = Counter(data[target]).most_common(1)[0][0]
        else:
            tree[best_feature][value] = id3(subset, features, target)
    return tree


# 打印决策树
def print_tree(tree, indent=""):
    if isinstance(tree, dict):
        for key, subtree in tree.items():
            print(indent + str(key))
            print_tree(subtree, indent + "  ")
    else:
        print(indent + str(tree))


# 加载数据
url = "https://archive.ics.uci.edu/ml/machine-learning-databases/wine/wine.data"
columns = ['Class', 'Alcohol', 'Malic acid', 'Ash', 'Alcalinity of ash', 'Magnesium', 'Total phenols', 'Flavanoids',
           'Nonflavanoid phenols', 'Proanthocyanins', 'Color intensity', 'Hue', 'OD280/OD315 of diluted wines',
           'Proline']
data = pd.read_csv(url, header=None, names=columns)

# 构建决策树
features = columns[1:]
target = columns[0]
tree = id3(data, features, target)

# 打印决策树
print("Decision Tree:")
print_tree(tree)
